
import networkx as nx
import time
#import matplotlib.pyplot as plt

def IskPartite(G):
    # check inner layer
    inner = filter(lambda x: G.node[x[0]]["layer"] == G.node[x[1]]["layer"], G.edges())
    if len(inner) > 0:
        return False
    
    # check cross layer
    for node in G.nodes():
        for nei in G.neighbors(node):
            lst = [G.node[node]["layer"], G.node[nei]["layer"]]
            if lst[0] - lst[1] > 1:
                return False
    
    return True


def kPartite(G, nlayer):
    # step1: eliminate cross layer links by adding virtual link
    # find all cross layer links
    for node in G.nodes():
        for nei in G.neighbors(node):
            lst = [G.node[node]["layer"], G.node[nei]["layer"]]
            if lst[0] - lst[1] > 1:
                G.remove_edge(node, nei)
                # calculate the node to be added
                path = [node]
                for t in range(lst[1]+1, lst[0]):
                    nname = str(node)+str(nei)+str(t)
                    G.add_node(nname, layer=t)
                    path.append(nname)
                path.append(nei)
                print path
                G.add_path(path)

    # step2: eliminate inner layer links
    # get connected component
    while (nlayer > 0):
        print "current is " + str(nlayer)
        # current layer
        GC = G.subgraph(filter(lambda x: G.node[x]["layer"] == nlayer, G.nodes()))
        
        if len(GC.edges()) == 0:
            print "skip layer"
            nlayer = nlayer - 1
            continue
        
        connset =  nx.connected_components(GC)
        # for each list, get a set with neighbor
        for l in connset:
            print 'Connected Set : ' + str(l)
            nlc = set()
            for n in l:
                nlc.update(G.neighbors(n))
            nlc.difference_update(l)
            print nlc
            #add edge
            for n in l:
                for s in nlc:
                    print "adding "+ str(n)+ ','+str(s)
                    G.add_edge(n, s)
        
        nlayer = nlayer - 1
    
    # remove those inner edges
    inner = filter(lambda x: G.node[x[0]]["layer"] == G.node[x[1]]["layer"], G.edges())
    print  'to be removed : ' + str(inner)
    for re in inner:
        G.remove_edge(re[0], re[1])
    


def main():
    
    edgelist=[('a','b'),('a','i'),('d','h'),('h','q'),('i','m')]
    edgelist.extend([('b','e'),('b','c'),('d','c'),('d','g'),('h','g'),('q','p'),('m','n'),('i','j')])
    edgelist.extend([('e','j'),('e','f'),('c','f'),('g','l'),('g','f'),('n','j'),('n','o'),('p','l'),('p','o')])
    edgelist.extend([('j','k'),('l','k'),('o','k'),('f','k')])
    
    G=nx.Graph()
    # add node with attribute layer
    l3 = ['a','b','d','h','q','i','m']    
    G.add_nodes_from(l3, layer=3)
    l2 = ['e','c','g','n','p']    
    G.add_nodes_from(l2, layer=2)
    l1 = ['j','l','o','f']    
    G.add_nodes_from(l1, layer=1)
    l0 = ['k']    
    G.add_nodes_from(l0, layer=0)
    ''' add edge without attribute. In fact, it should have attribute like length'''
    G.add_edges_from(edgelist)
    
    print "edge count: " + str(len(G.edges())) 
    
    kPartite(G,3)
                
    print "edge count: " + str(len(G.edges()))
    
    print IskPartite(G)
                
    nx.draw(G,pos=nx.spring_layout(G))
    #plt.show()

if __name__ == "__main__":
    start = time.clock()
    main()
    elapsed = (time.clock() - start)
    print "Time used: %.2f" % elapsed

